#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
@author: wang_chao03
@project: pyraft
@file: server
@time: 2023/2/17
"""

import asyncio
from typing import Optional, Union, Dict, List, Tuple, Iterator, Any, Set, Callable

from pyraft.state import State, Follower, Candidate, Leader
from pyraft.network import UDPProtocol
from pyraft.storage import StateMachine
from pyraft.log import logger


class Server:
    servers = set()

    def __init__(self, host: str, port: int, loop: Optional[asyncio.AbstractEventLoop] = None):
        self.host = host
        self.port = port
        self.loop = loop or asyncio.get_event_loop()
        self.cluster: Set[Tuple[str, int]] = set()
        self.requests = asyncio.Queue(loop=self.loop)
        self.state: Optional[State] = None
        self.transport: Optional[asyncio.DatagramTransport] = None
        self.__class__.servers.add(self)

    async def create_endpoint(self, protocol: UDPProtocol):
        if not self.transport:
            transport, _ = await self.loop.create_datagram_endpoint(
                protocol, local_addr=(self.host, self.port)
            )
            return transport
        return self.transport

    def udp_conn_lost_callback(self):
        ...

    async def start(self):
        protocol = UDPProtocol(
            queue=self.requests,
            request_handler=self.request_handler,
            conn_lost_callback=self.udp_conn_lost_callback,
            loop=self.loop
        )
        self.transport = await self.create_endpoint(protocol)
        self.state = State(self)
        self.state.start()

    def stop(self):
        self.state.stop()
        self.transport.close()

    def update_cluster(self, host: str, port: int):
        self.cluster.add((host, port))

    def request_handler(self, data: Dict, sender: Tuple[str, int]) -> None:
        self.state.request_handler(data, sender)

    def send(self, data: Any, dest: Union[str, Tuple[str, int]]):
        if isinstance(dest, str):
            host, port = dest.split(':')
            dest = (host, int(port))
        self.requests.put_nowait((data, dest))

    def broadcast(self, data: Any):
        for dest in self.cluster:
            self.send(data, dest)

    @staticmethod
    def add_state_apply_handler(handler: Optional[Callable[[StateMachine, Dict[str, Any]], None]]):
        State.add_state_apply_handler(handler)

    @staticmethod
    def add_follower_listener(callback: Callable[['Follower'], None]):
        State.add_follower_listener(callback)

    @staticmethod
    def add_candidate_listener(callback: Callable[['Candidate'], None]):
        State.add_candidate_listener(callback)

    @staticmethod
    def add_leader_listener(callback: Callable[['Leader'], None]):
        State.add_leader_listener(callback)


def parser_server_str(
        servers: Union[str, Iterator[str], Iterator[Tuple[str, int]]]
) -> List[Tuple[str, int]]:
    server_list = []
    if isinstance(servers, str):
        for server in servers.split(','):
            host, port = server.rsplit(':', 1)
            server_list.append((host, int(port)))
    elif isinstance(servers, (list, tuple)):
        for server in servers:
            if isinstance(server, str):
                host, port = server.rsplit(':', 1)
                server_list.append((host, int(port)))
            elif isinstance(server, (list, tuple)):
                host, port = server
                server_list.append((host, int(port)))
    return server_list


def leader_listener(role: Leader):
    logger.debug(f'当前服务切换为leader，ID为{role.id}')


def follower_listener(role: Follower):
    logger.debug(f'当前服务切换为follower，ID为{role.id}')


def candidate_listener(role: Candidate):
    logger.debug(f'当前服务切换为candidate，ID为{role.id}')


def state_apply_handler(state_machine: StateMachine, command: Dict[str, Any]):
    logger.debug(f'已将命令[{command}]应用到状态机')


async def start(
        servers: Union[str, Iterator[str], Iterator[Tuple[str, int]]],
        current_server_index: int,
        loop: Optional[asyncio.AbstractEventLoop] = None
) -> Server:
    loop = loop or asyncio.get_event_loop()
    servers = parser_server_str(servers)
    if current_server_index >= len(servers):
        raise IndexError(f'设置的当前服务器索引参数current_server_index[{current_server_index}]越界')
    current_host, current_port = servers.pop(current_server_index)
    current_server = Server(current_host, current_port, loop=loop)
    for host, port in servers:
        current_server.update_cluster(host, port)
    current_server.add_leader_listener(leader_listener)
    current_server.add_follower_listener(follower_listener)
    current_server.add_candidate_listener(candidate_listener)
    current_server.add_state_apply_handler(state_apply_handler)
    await current_server.start()
    logger.info(f'服务器[{current_host}:{current_port}]已启动')
    return current_server


def stop():
    for server in Server.servers:
        server.stop()
